import os
import sys
from embeddings.BaseEmbedding import BaseEmbedding
from sentence_transformers import SentenceTransformer 

# 由于只在文本嵌入使用到了GPU，所以直接指定第二块GPT
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

BGE_EMBEDDING_DIMS = 1024

EMBENDDING_MODEL = os.getenv("EMBEDDING_MODEL") 

if EMBENDDING_MODEL == None:
    EMBENDDING_MODEL = "bge-large-zh-v1.5"
    
ROOT_PATH = os.getenv("ROOT_PATH")
if ROOT_PATH == None:
    ROOT_PATH = sys.path[0]


class BgeEmbedding(BaseEmbedding):
    def __init__(self):
        super().__init__()
        self.model_path = os.path.join(ROOT_PATH, "model")
        self.model_path = os.path.join(self.model_path, EMBENDDING_MODEL)
        self.model = SentenceTransformer(self.model_path)

    #default dim
    @staticmethod
    def getDims():
        return BGE_EMBEDDING_DIMS
    
    # def __call__(self, strs: str):
    #     return self.getEmbedding(strs)
    def getEmbedding(self, strs: str):
        return self.model.encode(strs, normalize_embeddings=True)
    